from autokeras.image_supervised import load_image_dataset, ImageClassifier
from keras.models import load_model
from keras.utils import plot_model
from keras.preprocessing.image import load_img, img_to_array
import numpy as np

TRAIN_IMG_DIR = 'G:/doc7/Dogs vs Cats Redux Kernels Edition/train'
TRAIN_IMG_OUT_DIR = 'G:/doc7/Dogs vs Cats Redux Kernels Edition/train_out'
TRAIN_CSV_DIR = 'G:/doc7/Dogs vs Cats Redux Kernels Edition/train_labels.csv'
TEST_IMG_DIR = 'G:/doc7/Dogs vs Cats Redux Kernels Edition/test'

MODEL_DIR = 'G:/doc7/Dogs vs Cats Redux Kernels Edition/model/model.h5'
MODEL_PNG = 'G:/doc7/Dogs vs Cats Redux Kernels Edition/model/model.png'
IMAGE_SIZE = 28
 
if __name__ == '__main__':
    # 获取本地图片，转换成numpy格式
    train_data, train_labels = load_image_dataset(csv_file_path=TRAIN_CSV_DIR, images_path=TRAIN_IMG_OUT_DIR)

    # 数据进行格式转换
    train_data = train_data.astype('float32') / 255

    # 使用图片识别器
    clf = ImageClassifier(verbose=True)
    # 给其训练数据和标签，训练的最长时间可以设定，假设为1分钟，autokers会不断找寻最优的网络模型
    clf.fit(train_data, train_labels, time_limit=1 * 60)
    # 找到最优模型后，再最后进行一次训练和验证
    # 给出评估结果
    y = clf.evaluate(train_data, train_labels)
    print("evaluate:", y)

    # 导出我们生成的模型
    clf.load_searcher().load_best_model().produce_keras_model().save(MODEL_DIR)
    # 加载模型
    model = load_model(MODEL_DIR)
    # 将模型导出成可视化图片
    plot_model(model, to_file=MODEL_PNG)